import numpy

from histogram import Histogram1D, Histogram2D, Histogram3D

def quiet_sqrt(v):
    return numpy.where(v > 0, numpy.sqrt(numpy.fabs(v)), -numpy.sqrt(numpy.fabs(v)))


mumupi0pi0 = {'m_mumu_v_mm_off_pi0pi0': Histogram2D(name='m_mumu_v_mm_off_pi0pi0', 
                                              nbins=(100,100), 
                                              aranges=((9.35, 10.75), (9.35, 10.75)), 
                                              xlabel=r'M($\mu^+\mu^-$)', 
                                              ylabel=r'MM($\pi^0\pi^0$)', 
                                              xunits='GeV', 
                                              yunits='GeV',
                                              assoc_cols=('m_mumu', 'mm_off_pi0pi0'),
                                              ),
              'r2_v_m_mumu': Histogram2D(name='r2_v_m_mumu', 
                                         nbins=(100,200), 
                                         aranges=((0,1), (9, 11)), 
                                         xlabel=r'R2', 
                                         ylabel=r'M($\mu^+\mu^-$)', 
                                         xunits='', 
                                         yunits='GeV',
                                         assoc_cols=('r2', 'm_mumu'),
                                         ),
              'mbc_mumupi0pi0': Histogram1D(name='mbc_mumupi0pi0',
                                            nbins=100,
                                            arange=(-0.5, 0.5),
                                            xlabel=r'$M_{BC}(\mu^+\mu^-\pi^0\pi^0)$',
                                            xunits=r'GeV',
                                            assoc_cols='mbc_mumupi0pi0',
                                            ),
              }

mumu = {'m_mumu': Histogram1D(name='m_mumu',
                             nbins=140,
                             arange=(9.35, 10.75),
                             xlabel=r'M($\mu^+\mu^-$)',
                             xunits='GeV',
                             assoc_cols='m_mumu',
                             ),
        'm_mumu_20MeVbins': Histogram1D(name='m_mumu_20MeVbins',
                                        nbins=70,
                                        arange=(9.35, 10.75),
                                        xlabel=r'M($\mu^+\mu^-$)',
                                        xunits='GeV',
                                        assoc_cols='m_mumu',
                                        ),
        }

pi0pi0 = { 'm2_off_pi0_dalitz': Histogram2D(name='m2_off_pi0_dalitz',
                                            nbins=(320, 320),
                                            aranges=((100, 116), (100, 116)),
                                            xlabel=r'MM($\pi^0_1$)',
                                            ylabel=r'MM($\pi^0_2$)',
                                            xunits='GeV$^2$',
                                            yunits='GeV$^2$',
                                            assoc_cols=('m2_off_pi01',
                                                        'm2_off_pi02'),
                                            ),
           'mm_off_pi01': Histogram1D(name='mm_off_pi01',
                                      nbins=200,
                                      arange=(9, 11),
                                      xlabel=r'MM($\pi^0_1$)',
                                      xunits=r'GeV',
                                      assoc_cols='m2_off_pi01',
                                      ),
           'mm_off_pi02': Histogram1D(name='mm_off_pi02',
                                      nbins=200,
                                      arange=(9, 11),
                                      xlabel=r'MM($\pi^0_2$)',
                                      xunits=r'GeV',
                                      assoc_cols='m2_off_pi02',
                                      ),
            'mm_off_pi0pi0': Histogram1D(name='mm_off_pi0pi0',
                                         nbins=320,
                                         arange=(9, 10.6),
                                         xlabel=r'MM($\pi^0\pi^0$)',
                                         xunits='GeV',
                                         assoc_cols='mm_off_pi0pi0',
                                        ),
           'm_pi0pi0':  Histogram1D(name='m_pi0pi0',
                                        nbins=340,
                                        arange=(0.15, 1.85),
                                        xlabel=r'M($\pi^0\pi^0$)',
                                        xunits='GeV',
                                        assoc_cols='m_pi0pi0',
                                        ),
           'mm_off_pi0pi0_20MeVbins': Histogram1D(name='mm_off_pi0pi0_20MeVbins',
                                                  nbins=70,
                                                  arange=(9.35, 10.75),
                                                  xlabel=r'MM($\pi^0\pi^0$)',
                                                  xunits='GeV',
                                                  assoc_cols='mm_off_pi0pi0',
                                                  ),       
           'pi01_g1e': Histogram1D(name='pi01_g1e',
                                   nbins=1000,
                                   arange=(0, 1),
                                   xlabel=r'$\pi^0_1 E\left[\gamma_1\right]$',
                                   xunits='GeV',
                                   assoc_cols='pi01_g1e',
                                   ),
           'pi01_g2e': Histogram1D(name='pi01_g2e',
                                   nbins=1000,
                                   arange=(0, 1),
                                   xlabel=r'$\pi^0_1 E\left[\gamma_2\right]$',
                                   xunits='GeV',
                                   assoc_cols='pi01_g2e',
                                   ), 
           'pi02_g1e': Histogram1D(name='pi02_g1e',
                                   nbins=1000,
                                   arange=(0, 1),
                                   xlabel=r'$\pi^0_2 E\left[\gamma_1\right]$',
                                   xunits='GeV',
                                   assoc_cols='pi02_g1e',
                                   ),
           'pi02_g2e': Histogram1D(name='pi02_g2e',
                                   nbins=1000,
                                   arange=(0, 1),
                                   xlabel=r'$\pi^0_2 E\left[\gamma_2\right]$',
                                   xunits='GeV',
                                   assoc_cols='pi02_g2e',
                                   ), 
           
           }



hist_colfilters = {'mm_off_pi01': quiet_sqrt,
                  'mm_off_pi02': quiet_sqrt,
                 }


hists = dict(**mumupi0pi0)
hists.update(mumu)
hists.update(pi0pi0)


histsets = {'default': hists.keys(),
            'mumu': mumu.keys(),
            'pi0pi0': pi0pi0.keys(),
            }


    
